import requests
import json
from env import ERNIE_SPEED_API_KEY, ERNIE_SPEED_SECRET_KEY


class ErnieSpeed:

    def get_access_token(self):
        """
        使用 API Key，Secret Key 获取access_token，替换下列示例中的应用API Key、应用Secret Key
        """
        app_id = 88840877
        api_key = ERNIE_SPEED_API_KEY
        secret_key = ERNIE_SPEED_SECRET_KEY
        url = (
            "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=%s&client_secret=%s"
            % (api_key, secret_key)
        )

        payload = json.dumps("")
        headers = {"Content-Type": "application/json", "Accept": "application/json"}

        response = requests.request("POST", url, headers=headers, data=payload)
        return response.json().get("access_token")

    def chat(self, messages):

        access_token = self.get_access_token()
        # print(access_token)
        url = (
            "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-speed-128k?access_token="
            + access_token
        )

        payload = json.dumps({"messages": messages})
        headers = {"Content-Type": "application/json"}

        response = requests.request("POST", url, headers=headers, data=payload)
        resp = json.loads(response.text)
        print(resp)
        return resp
    
    def getMessage(self, resp):
        return resp.get("result", "");
